import { create } from 'zustand'
import { Message } from 'ai/react'

import { avatars, type Avatar } from '@/mock/avatars'

type AvatarState = {
    avatar: Avatar
    setAvatar: (avatar: Avatar) => void

    messageGroup: Record<string, Message[]>
    messagesWith: (avatarId: string) => Message[]
    updateMessages: (avatarId: string, messages: Message[]) => void
}

export const useAvatar = create<AvatarState>((set, get) => ({
    avatar: avatars[0],
    setAvatar: (avatar: Avatar) => set(state => ({ avatar: avatar })),

    messageGroup: {},
    messagesWith: (avatarId: string) => get().messageGroup[avatarId],
    updateMessages: (avatarId: string, messages: Message[]) => {
        const record: Record<string, Message[]> = {}
        record[avatarId] = messages
        set(state => ({ messageGroup: Object.assign(get().messageGroup, record) }))
    }
}))